import math
import numpy as np
import pandas as pd
from multiprocessing import Process, Queue


np.set_printoptions(formatter={'float':'{:g}'.format}) #specifies numpy array number format

print('Loading Functions...')

#Constants and Unit Conversion
G=6.67408*10**(-11) #m^3 kg^-1 s^-2
Msun=1.988*10**30 #kg
u=Msun*G
MEarth=5.9722*10**24 #kg
ue=MEarth*G
REarth=149.6*10**9 #meters
ParkAlt=100000 #Earth Parking Orbit altitude above planet surface
VEarth=math.sqrt(u*(float(2)/float(REarth)-float(1)/float(REarth)))
Vpark=math.sqrt(ue*(float(2)/float(6378000+ParkAlt)-float(1)/float(6378000+ParkAlt))) #orbital velocity at parking orbit
Vesc=math.sqrt(2*ue/(6378000+ParkAlt)) #escape velocity at parking orbit
def MtoAU(m):
    return m*6.68459*10**(-12)
def AUtoM(AU):
    return AU/(6.68459*10**(-12))
def DegtoRad(deg):
    return deg*math.pi/float(180)
def RadtoDeg(rad):
    return rad*180/math.pi

#Final Functions
def DeltaV3(w,i,ec,a):
    Total=M13(ec,AUtoM(a))+M23(DegtoRad(w),DegtoRad(i),ec,AUtoM(a))+M3(DegtoRad(w),DegtoRad(i),ec,AUtoM(a))
    return Total
def DeltaV3exp(w,i,ec,a):
    return M13(ec,AUtoM(a)),M23(DegtoRad(w),DegtoRad(i),ec,AUtoM(a)),M3(DegtoRad(w),DegtoRad(i),ec,AUtoM(a))
def DVArray3(array):
    return DeltaV3(array[0],array[1],array[2],array[3])

def DeltaV2(w,i,ec,a):
    return M12(DegtoRad(w),ec,AUtoM(a))+M22(DegtoRad(w),DegtoRad(i),ec,AUtoM(a))
def DeltaVexp2(w,i,ec,a):
    return M12(DegtoRad(w),ec,AUtoM(a)),M22(DegtoRad(w),DegtoRad(i),ec,AUtoM(a))
def DVArray2(array):
    return DeltaV2(array[0],array[1],array[2],array[3])

def Travel3(ec,a):
    return 1/2*math.sqrt(4*math.pi**2/u*TransA3(ec,AUtoM(a))**3)/86400
def Travel3Array(array):
    return Travel3(array[2],array[3])

def Travel2(w,ec,a):
    return 1/2*math.sqrt(4*math.pi**2/u*((RNode(DegtoRad(w),ec,AUtoM(a))+REarth)/float(2))**3)/86400
def Travel2Array(array):
    return Travel2(array[0],array[2],array[3])


#Maneuver 1, 3 Burn
#Start at 100 km Earth parking orbit
def AstAp(ec,a): #apoapsis of asteroid
    return (1+ec)*a
def AstPe(ec,a): #periapsis of asteroid
    return (1-ec)*a
def velo(semi,r): #orbital velocity at a radius
    return math.sqrt(u*(float(2)/float(r)-float(1)/float(semi)))
def TransPeV3(ec,a): #velocity at transfer periapsis
    return velo((AstAp(ec,a)+REarth)/float(2),REarth)
def M13(ec,a): #Delta-V for Maneuver 1
    return math.sqrt(((TransPeV3(ec,a)-VEarth)**2)+(Vesc**2))-Vpark

#Maneuver 2, 3 Burn
def TransA3(ec,a):
    return (AstAp(ec,a)+REarth)/float(2)
def TransEC3(ec,a): #Transfer orbit Eccentricity
    return float((AstAp(ec,a)-REarth))/(AstAp(ec,a)+REarth)
def perilat(w,i): #"Latitude" of Periapsis
    return math.asin(math.sin(w)*math.sin(i))
def TransR3(ec,a): #Radius at Maneuver 2 burn
    return TransA3(ec,a)*(1-TransEC3(ec,a)**2)
def FPAngle(ec,a): #Flight Path Angle
    return math.acos(1/(math.sqrt(1+TransEC3(ec,a)**2)))
def M23Angle(w,i,ec,a): #Maneuver 2 Angle
    return math.acos(math.cos(perilat(w,i))*(math.cos(FPAngle(ec,a))**2)+(math.sin(FPAngle(ec,a))**2))
def M23(w,i,ec,a): #Delta-V for Maneuver 2, derived from law of cosines with equal side lengths
    return 2*velo(TransA3(ec,a),TransR3(ec,a))*math.sin(M23Angle(w,i,ec,a)/2)

#Maneuver 3, 3 Burn
def TwistAng(w,i): #angle for very last maneuver
    return math.acos(math.cos(i)/math.sqrt((math.cos(w)**2)+((math.cos(i)**2)*(math.sin(w)**2))))
def M3Vi(ec,a): #initial velocity at Maneuver 3 position
    return velo(TransA3(ec,a),AstAp(ec,a))
def M3Vf(ec,a): #final velocity at Maneuver 3 position
    return velo(a,AstAp(ec,a))
def M3(w,i,ec,a): #Delta-V for Maneuver 3, derived from law of cosines
    return math.sqrt(M3Vi(ec,a)**2+M3Vf(ec,a)**2-2*M3Vi(ec,a)*M3Vf(ec,a)*math.cos(TwistAng(w,i)))

#Burn 1, 2 Burn
def RNode(w,ec,a): #radius at ascending or decending node, choose higher one
    Node1=a*(1-(ec**2))/(1+ec*math.cos(-w))
    Node2=a*(1-(ec**2))/(1+ec*math.cos(math.pi-w))
    if Node1>Node2:
        return Node1
    else:
        return Node2
def NodeChoice(w,ec,a): #which node was chosen, for all future calculation the periapsis direction is defined as 0 degrees
    Node1=a*(1-(ec**2))/(1+ec*math.cos(-w))
    Node2=a*(1-(ec**2))/(1+ec*math.cos(math.pi-w))
    if Node1>Node2:
        return -w
    else:
        return math.pi-w
def TransPeV2(w,ec,a): #velocity at transfer periapsis
    return velo((RNode(w,ec,a)+REarth)/float(2),REarth)
def M12(w,ec,a): #Delta-V for Maneuver 1
    return math.sqrt(((TransPeV2(w,ec,a)-VEarth)**2)+(Vesc**2))-Vpark

#Burn 2, 2 Burn
#Define Orbital Elements
def TransA2(w,ec,a):
    return (RNode(w,ec,a)+REarth)/float(2)
def TransEC2(w,ec,a): #Transfer orbit Eccentricity
    return float((RNode(w,ec,a)-REarth))/(RNode(w,ec,a)+REarth)
def AstB(ec,a): #semiminor axis of asteroid orbit
    return a*math.sqrt(1-ec**2)
def AstC(ec,a): #"c" of ellipse, useful for calculation
    return a*ec
#Define Velocity magnitudes
def TransVM2(w,ec,a):
    return velo(TransA2(w,ec,a),TransA2(w,ec,a)*(1+TransEC2(w,ec,a)))
def AstVM2(w,ec,a):
    return velo(a,RNode(w,ec,a))
#Define Vector components
def VecX(w,ec,a): #X direction vector component of asteroid orbit at node
    return -RNode(w,ec,a)*math.sin(NodeChoice(w,ec,a))/AstB(ec,a)**2
def VecY(w,ec,a): #Y direction vector component of asteroid orbit at node
    return (RNode(w,ec,a)*math.cos(NodeChoice(w,ec,a))+AstC(ec,a))/a**2
def NormConstAst(w,ec,a):#Normalization constant for Velocity vector
    return math.sqrt(VecX(w,ec,a)**2+VecY(w,ec,a)**2)
def VecXTrans(w,i,ec,a): #X direction vector component of transfer orbit at node
    return -math.sin(NodeChoice(w,ec,a))*math.cos(i)
def VecYTrans(w,i,ec,a): #Y direction vector component of transfer orbit at node
    return math.cos(NodeChoice(w,ec,a))*math.cos(i)
def VecZTrans(i):
    return math.sin(i)
def TransNorm(w,i,ec,a):
    return math.sqrt(VecXTrans(w,i,ec,a)**2+VecYTrans(w,i,ec,a)**2+VecZTrans(i)**2)
def FinalM2Angle(w,i,ec,a): #Calculate Angle between Orbits using Dot Product
    return math.acos((VecX(w,ec,a)*VecXTrans(w,i,ec,a)+(VecY(w,ec,a)*VecYTrans(w,i,ec,a)))/(NormConstAst(w,ec,a)*TransNorm(w,i,ec,a)))
def M22(w,i,ec,a): #Delta-V for Maneuver 1, derived from law of cosines
    return math.sqrt(TransVM2(w,ec,a)**2+AstVM2(w,ec,a)**2-2*TransVM2(w,ec,a)*AstVM2(w,ec,a)*math.cos(FinalM2Angle(w,i,ec,a)))

print('Functions Loaded')

print('Downloading Database...')

import urllib.request
# Download the file from `url` and save it locally under `file_name`:
urllib.request.urlretrieve('http://www.minorplanetcenter.net/iau/MPCORB/MPCORB.DAT.gz', 'MPCORB.DAT.gz')

print('Database Downloaded')

print('Importing Data...')
#Load datafile, line 44 is data, skip_header=43
RAWDATA=np.genfromtxt('MPCORB.DAT.gz',delimiter=(7,5,7,6,10,11,11,11,11,12,12,2,1,10,6,4,10,6,3,5),skip_header=43,skip_footer=0,usecols=(5,7,8,10),dtype="f16")
RAWNAMES=np.genfromtxt('MPCORB.DAT.gz',delimiter=(7,5,7,6,10,11,11,11,11,12,12,2,1,10,6,4,10,6,3,5),skip_header=43,skip_footer=0,usecols=(0,12,18),dtype="str")

ToPandas=np.rec.fromarrays((RAWNAMES[:,0],RAWNAMES[:,1],RAWNAMES[:,2],RAWDATA[:,0],RAWDATA[:,1],RAWDATA[:,2],RAWDATA[:,3]),names=('Names','U','Pert','W', 'I','EC','A'))
df=pd.DataFrame(ToPandas)
df2=df[~df['Pert'].isin(['   '])]
df3=df2[~df2['U'].isin([' '])]
PandasOut=df3[~df3['Names'].isin(['\n'])]
NAMES=PandasOut.as_matrix(columns=PandasOut.columns[0:2])
DATA=PandasOut.as_matrix(columns=PandasOut.columns[3:7])
print('Import Complete')

print('Calculating Delta-Vs...')

def DeltaVCalc(DVArray,InputData,DVOut):
    DVOut.put(np.apply_along_axis(DVArray,1,InputData))

sims={DVArray3,DVArray2}
    
if __name__ == '__main__':
    DVOut=Queue()
    processes = [Process(target=DeltaVCalc, args=(b,DATA,DVOut,)) for b in sims]
    for p in processes:
        p.start()
    Results=[DVOut.get() for p in processes]
print('Delta-V Calculations Complete')

print('Calculating Travel Times...')
TravelTimes=np.asarray((np.apply_along_axis(Travel3Array,1,DATA),np.apply_along_axis(Travel2Array,1,DATA))).T
print('Travel Time Calculations Complete')

print('Creating Output Datafile...')
ResultTableTemp1=np.append(DATA.T,[Results[0]],axis=0).T
ResultTableTemp2=np.append(ResultTableTemp1.T,[TravelTimes[:,0]],axis=0).T
ResultTableTemp3=np.append(ResultTableTemp2.T,[Results[1]],axis=0).T
ResultTable=np.append(ResultTableTemp3.T,[TravelTimes[:,1]],axis=0).T

def best(array):
    if array[6]>array[4]:
        return array[4],array[5],3
    else:
        return array[6],array[7],2

MinResult=np.apply_along_axis(best,1,ResultTable)

FullResultTable=np.hstack((ResultTable,MinResult))

FinalDataTable=np.hstack((NAMES,FullResultTable))
MBFinalDataTable=FinalDataTable[(5.203>FinalDataTable[:,5]*(1+FinalDataTable[:,4])) & (FinalDataTable[:,5]*(1-FinalDataTable[:,4])>1.524)]

np.savetxt('DeltaV_Results.txt',FinalDataTable, 
           fmt=('%-7s','%-1s','%9.5f','%9.5f','%9.7f','%12.7f','%5.0f','%4.0f','%5.0f','%4.0f','%5.0f','%4.0f','%i'),newline='\r\n',
          header='Desn  U  Peri.      Incl.    e            a        3Burn  3TT 2Burn  2TT  Best   TT Burns')
np.savetxt('DeltaV_Results_MB.txt',MBFinalDataTable, 
           fmt=('%-7s','%-1s','%9.5f','%9.5f','%9.7f','%12.7f','%5.0f','%4.0f','%5.0f','%4.0f','%5.0f','%4.0f','%i'),newline='\r\n',
          header='Desn  U  Peri.      Incl.    e            a        3Burn  3TT 2Burn  2TT  Best   TT Burns')


print('Output Datafile Created', '\nAll Processes Complete')
